import numpy as np
import cv2
import os
from config import cfg
from tqdm import tqdm
from torch.utils.data import Dataset
from PIL import Image
import torchvision
import time

def walk(path):
    file_list = []
    ff = os.walk(path)
    for root, dirs, files in ff:
        for file in files:
            file_list.append(os.path.join(root, file))
    return file_list


def compute_meanANDstd():
    # img_h, img_w = 32, 32
    img_h, img_w = cfg.DATA.IMG_OUT_WIDTH, cfg.DATA.IMG_OUT_HEIGHT  # 根据自己数据集适当调整，影响不大
    means, stdevs = [], []

    imgs_path = cfg.DATA.OUT_PATH
    img_list = walk(imgs_path)
    len_ = len(img_list)

    img_sets = []
    i = 0
    for item in tqdm(img_list):
        #img = cv2.imread(os.path.join(imgs_path, item))

        img = cv2.imread(item)
        img = cv2.resize(img, (img_w, img_h))
        img = img[:, :, :, np.newaxis]
        img_sets.append(img)
        i += 1
        #print(i, '/', len_)

    imgs = np.concatenate(img_sets, axis=3)
    imgs = imgs.astype(np.float32) / 255.

    for i in range(3):
        pixels = imgs[:, :, i, :].ravel()  # 拉成一行
        means.append(np.mean(pixels))
        stdevs.append(np.std(pixels))

    # BGR --> RGB ， CV读取的需要转换，PIL读取的不用转换
    means.reverse()
    stdevs.reverse()

    print("normMean = {}".format(means))
    print("normStd = {}".format(stdevs))




def compute_mean_and_std(dataset):
    # 输入PyTorch的dataset，输出均值和标准差
    mean_r = 0
    mean_g = 0
    mean_b = 0
    print("计算均值>>>")
    for img_path, _ in tqdm(dataset,ncols=80):
        img=Image.open(img_path)
        img = np.asarray(img) # change PIL Image to numpy array
        mean_b += np.mean(img[:, :, 0])
        mean_g += np.mean(img[:, :, 1])
        mean_r += np.mean(img[:, :, 2])

    mean_b /= len(dataset)
    mean_g /= len(dataset)
    mean_r /= len(dataset)

    diff_r = 0
    diff_g = 0
    diff_b = 0

    N = 0
    print("计算方差>>>")
    for img_path, _ in tqdm(dataset,ncols=80):
        img=Image.open(img_path)
        img = np.asarray(img)
        diff_b += np.sum(np.power(img[:, :, 0] - mean_b, 2))
        diff_g += np.sum(np.power(img[:, :, 1] - mean_g, 2))
        diff_r += np.sum(np.power(img[:, :, 2] - mean_r, 2))

        N += np.prod(img[:, :, 0].shape)

    std_b = np.sqrt(diff_b / N)
    std_g = np.sqrt(diff_g / N)
    std_r = np.sqrt(diff_r / N)

    mean = (mean_b.item() / 255.0, mean_g.item() / 255.0, mean_r.item() / 255.0)
    std = (std_b.item() / 255.0, std_g.item() / 255.0, std_r.item() / 255.0)
    return mean, std

def main():
    path = cfg.IMGLIST.IMG_OUT_DIR
    train_path = os.path.join(path,'train_data')
    val_path = os.path.join(path,'val_data')
    #test_path = path + '/test'

    train_data = torchvision.datasets.ImageFolder(train_path)
    val_data = torchvision.datasets.ImageFolder(val_path)
    #test_data = torchvision.datasets.ImageFolder(test_path)

    # train_mean,train_std=compute_mean_and_std(train_data.imgs)
    time_start = time.time()
    train_mean, train_std = compute_mean_and_std(train_data.imgs)
    time_end = time.time()
    print()
    print("训练集计算消耗时间：", round(time_end - time_start, 4), "s")

    time_start = time.time()
    val_mean, val_std = compute_mean_and_std(val_data.imgs)
    time_end = time.time()
    print("验证集计算消耗时间：", round(time_end - time_start, 4), "s")

    # test_mean,test_std=compute_mean_and_std(test_data.imgs)
    print("训练集的平均值：{}".format(train_mean))
    print("训练集的方差：{}".format(train_std))

    print("验证集的平均值：{}".format(val_mean))
    print("验证集的方差：{}".format(val_std))


if __name__ == '__main__':
    main()
